{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Title: #Maximum Balanced Subsequence Sum"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Difficulty: #Hard"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Category Title: #Algorithms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Tag Slug: #binary-indexed-tree #segment-tree #array #binary-search #dynamic-programming"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Name Translated: #树状数组 #线段树 #数组 #二分查找 #动态规划"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solution Name: maxBalancedSubsequenceSum"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Translated Title: #平衡子序列的最大和"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Translated Content:\n",
    "<p>给你一个下标从 <strong>0</strong>&nbsp;开始的整数数组&nbsp;<code>nums</code>&nbsp;。</p>\n",
    "\n",
    "<p><code>nums</code>&nbsp;一个长度为 <code>k</code>&nbsp;的 <strong>子序列</strong>&nbsp;指的是选出 <code>k</code>&nbsp;个 <strong>下标</strong>&nbsp;<code>i<sub>0</sub>&nbsp;&lt;&nbsp;i<sub>1</sub> &lt;&nbsp;... &lt; i<sub>k-1</sub></code>&nbsp;，如果这个子序列满足以下条件，我们说它是 <strong>平衡的</strong>&nbsp;：</p>\n",
    "\n",
    "<ul>\n",
    "\t<li>对于范围&nbsp;<code>[1, k - 1]</code>&nbsp;内的所有&nbsp;<code>j</code>&nbsp;，<code>nums[i<sub>j</sub>] - nums[i<sub>j-1</sub>] &gt;= i<sub>j</sub> - i<sub>j-1</sub></code>&nbsp;都成立。</li>\n",
    "</ul>\n",
    "\n",
    "<p><code>nums</code>&nbsp;长度为 <code>1</code>&nbsp;的 <strong>子序列</strong>&nbsp;是平衡的。</p>\n",
    "\n",
    "<p>请你返回一个整数，表示 <code>nums</code>&nbsp;<strong>平衡</strong>&nbsp;子序列里面的 <strong>最大元素和</strong>&nbsp;。</p>\n",
    "\n",
    "<p>一个数组的 <strong>子序列</strong>&nbsp;指的是从原数组中删除一些元素（<strong>也可能一个元素也不删除</strong>）后，剩余元素保持相对顺序得到的 <strong>非空</strong>&nbsp;新数组。</p>\n",
    "\n",
    "<p>&nbsp;</p>\n",
    "\n",
    "<p><strong class=\"example\">示例 1：</strong></p>\n",
    "\n",
    "<pre>\n",
    "<b>输入：</b>nums = [3,3,5,6]\n",
    "<b>输出：</b>14\n",
    "<b>解释：</b>这个例子中，选择子序列 [3,5,6] ，下标为 0 ，2 和 3 的元素被选中。\n",
    "nums[2] - nums[0] &gt;= 2 - 0 。\n",
    "nums[3] - nums[2] &gt;= 3 - 2 。\n",
    "所以，这是一个平衡子序列，且它的和是所有平衡子序列里最大的。\n",
    "包含下标 1 ，2 和 3 的子序列也是一个平衡的子序列。\n",
    "最大平衡子序列和为 14 。</pre>\n",
    "\n",
    "<p><strong class=\"example\">示例 2：</strong></p>\n",
    "\n",
    "<pre>\n",
    "<b>输入：</b>nums = [5,-1,-3,8]\n",
    "<b>输出：</b>13\n",
    "<b>解释：</b>这个例子中，选择子序列 [5,8] ，下标为 0 和 3 的元素被选中。\n",
    "nums[3] - nums[0] &gt;= 3 - 0 。\n",
    "所以，这是一个平衡子序列，且它的和是所有平衡子序列里最大的。\n",
    "最大平衡子序列和为 13 。\n",
    "</pre>\n",
    "\n",
    "<p><strong class=\"example\">示例 3：</strong></p>\n",
    "\n",
    "<pre>\n",
    "<b>输入：</b>nums = [-2,-1]\n",
    "<b>输出：</b>-1\n",
    "<b>解释：</b>这个例子中，选择子序列 [-1] 。\n",
    "这是一个平衡子序列，而且它的和是 nums 所有平衡子序列里最大的。\n",
    "</pre>\n",
    "\n",
    "<p>&nbsp;</p>\n",
    "\n",
    "<p><strong>提示：</strong></p>\n",
    "\n",
    "<ul>\n",
    "\t<li><code>1 &lt;= nums.length &lt;= 10<sup>5</sup></code></li>\n",
    "\t<li><code>-10<sup>9</sup> &lt;= nums[i] &lt;= 10<sup>9</sup></code></li>\n",
    "</ul>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Description: [maximum-balanced-subsequence-sum](https://leetcode.cn/problems/maximum-balanced-subsequence-sum/description/)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solutions: [maximum-balanced-subsequence-sum](https://leetcode.cn/problems/maximum-balanced-subsequence-sum/solutions/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_cases = ['[3,3,5,6]', '[5,-1,-3,8]', '[-2,-1]']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class TreeNode:\n",
    "    def __init__(self, val=0, start=None, end=None, left=None, right=None):\n",
    "        self.val = val\n",
    "        self.start = start\n",
    "        self.end = end\n",
    "        self.left = left\n",
    "        self.right = right\n",
    "\n",
    "class SegmentTree:\n",
    "    def __init__(self, start, end, vals, mode='max'):\n",
    "        def build(start, end, vals):\n",
    "            if start == end:\n",
    "                return TreeNode(vals[start], start, end)\n",
    "            mid = (start + end) // 2\n",
    "            l, r = build(start, mid, vals), build(mid + 1, end, vals)\n",
    "            return TreeNode(max(l.val, r.val), start, end, l, r) if mode == 'max' else TreeNode(l.val+r.val, start, end, l, r)\n",
    "\n",
    "        self.t = build(start, end, vals)\n",
    "        self.mode = mode\n",
    "\n",
    "    def update(self, root, index, value):\n",
    "        if index == root.start == root.end:\n",
    "            root.val = value\n",
    "            return\n",
    "        mid = (root.start + root.end) // 2\n",
    "        if index <= mid:\n",
    "            self.update(root.left, index, value)\n",
    "        else:\n",
    "            self.update(root.right, index, value)\n",
    "        root.val = max(root.left.val, root.right.val) if self.mode == 'max' else root.left.val + root.right.val\n",
    "\n",
    "    def query(self, root, i, j):\n",
    "        if root.start == i and root.end == j:\n",
    "            return root.val\n",
    "        mid = (root.start + root.end) // 2\n",
    "        if j <= mid:\n",
    "            return self.query(root.left, i, j)\n",
    "        elif i > mid:\n",
    "            return self.query(root.right, i, j)\n",
    "        else:\n",
    "            if self.mode == 'max':\n",
    "                return max(self.query(root.left, i, mid), self.query(root.right, mid + 1, j))\n",
    "            else:\n",
    "                return self.query(root.left, i, mid) + self.query(root.right, mid + 1, j)\n",
    "\n",
    "class Solution:\n",
    "    def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:\n",
    "        \n",
    "        max_val = max(nums)       \n",
    "        if max_val <= 0:\n",
    "            return max_val\n",
    "        \n",
    "        n = len(nums)\n",
    "        \n",
    "        dp = [0] * n\n",
    "        \n",
    "        d = {k: i for i, k in enumerate(sorted(set(nums[x]-x for x in range(n))))}\n",
    "        \n",
    "        st = SegmentTree(0, len(d)-1, [0] * len(d))\n",
    "        \n",
    "        \n",
    "        for i in range(n):\n",
    "            res = st.query(st.t, 0, d[nums[i]-i])\n",
    "            dp[i] = max(nums[i], nums[i]+res)\n",
    "            if dp[i] > st.query(st.t, d[nums[i]-i], d[nums[i]-i]):\n",
    "                st.update(st.t, d[nums[i]-i], dp[i])\n",
    "                \n",
    "        \n",
    "        return max(dp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:\n",
    "        ans=max(nums)\n",
    "        if ans<=0:\n",
    "            return ans \n",
    "        n=len(nums)\n",
    "        delta=[nums[i]-i for i in range(n)]\n",
    "        delta=[delta[i] for i in range(n) if nums[i]>0]\n",
    "        nums=[i for i in nums if i>0]\n",
    "        n=len(delta)\n",
    "        keys=[delta[0]]\n",
    "        vals=[nums[0]]\n",
    "        #维护两个单增序列\n",
    "        for i in range(1,n):\n",
    "            b=bisect_left(keys,delta[i])\n",
    "            if b>=len(keys):\n",
    "                keys.append(delta[i])\n",
    "                vals.append(vals[-1]+nums[i])\n",
    "            elif keys[b]==delta[i]:\n",
    "                tmp=vals[b]+nums[i]\n",
    "                pos=bisect_right(vals,tmp)\n",
    "                vals[b]=tmp \n",
    "                keys=keys[:b+1]+keys[pos:]\n",
    "                vals=vals[:b+1]+vals[pos:]\n",
    "                \n",
    "            else:\n",
    "                if b==0:\n",
    "                    pos=bisect_right(vals,nums[i])\n",
    "                    keys=[delta[i]]+keys[pos:]\n",
    "                    vals=[nums[i]]+vals[pos:]\n",
    "                else:\n",
    "                    tmp=vals[b-1]+nums[i] \n",
    "                    pos=bisect_right(vals,tmp)\n",
    "                    keys=keys[:b]+[delta[i]]+keys[pos:]\n",
    "                    vals=vals[:b]+[tmp]+vals[pos:]\n",
    "        return vals[-1]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:\n",
    "        t = []\n",
    "        n = len(nums)\n",
    "        for i in range(n):\n",
    "            t.append(nums[i] - i)\n",
    "        t.sort()\n",
    "        tree = BIT(n)\n",
    "        for i, num in enumerate(nums):\n",
    "            rank = bisect_left(t, num - i) + 1\n",
    "            mx = max(tree.query(rank), 0) + num\n",
    "            tree.update(rank, mx)\n",
    "        return tree.query(n)\n",
    "\n",
    "class  BIT:\n",
    "    def __init__(self, n):\n",
    "        self.n = n\n",
    "        self.tree = [-inf] * (n + 1)\n",
    "\n",
    "    def update(self, i ,val):\n",
    "        while i <= self.n:\n",
    "            self.tree[i] = max(self.tree[i], val)\n",
    "            i += i & -i\n",
    "    def query(self, i) -> int:\n",
    "        mx = -inf\n",
    "        while i > 0:\n",
    "            mx = max(mx, self.tree[i])\n",
    "            i -= i & -i\n",
    "        return mx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:\n",
    "        # b = sorted(set(x-i for i, x in enumerate(nums)))\n",
    "        b = sorted(x-i for i, x in enumerate(nums))\n",
    "        print(f\"{b=}\")\n",
    "        t = BIT(len(b)+1)\n",
    "        ans = -inf\n",
    "        for i, x in enumerate(nums):\n",
    "            j = bisect_left(b, x-i) +1\n",
    "            f = max(t.pre_max(j), 0) + x\n",
    "            ans = max(ans, f)\n",
    "            t.update(j, f)\n",
    "        \n",
    "        return ans\n",
    "\n",
    "class BIT:\n",
    "    def __init__(self, n):\n",
    "        self.tree = [-inf] * n\n",
    "    \n",
    "    def update(self, i, val):\n",
    "        while i < len(self.tree):\n",
    "            self.tree[i] = max(self.tree[i], val)\n",
    "            i += i & -i\n",
    "    \n",
    "    def pre_max(self, i):\n",
    "        mx = -inf\n",
    "        while i > 0:\n",
    "            mx = max(self.tree[i], mx)\n",
    "            i &= i-1\n",
    "        return mx\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class BIT:\n",
    "    def __init__(self, n):\n",
    "        self.array = [- float('inf')] * (n + 1)\n",
    "    \n",
    "    def update(self, idx, val):\n",
    "        idx += 1 \n",
    "        while idx < len(self.array):\n",
    "            self.array[idx] = max(self.array[idx], val)\n",
    "            idx += idx & (- idx)\n",
    "    \n",
    "    def query(self, idx):\n",
    "        res = - float('inf')\n",
    "        idx += 1 \n",
    "        while idx > 0:\n",
    "            res = max(res, self.array[idx])\n",
    "            idx -= idx & (- idx)\n",
    "        return res \n",
    "\n",
    "class Solution:\n",
    "    def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:\n",
    "        n = len(nums)\n",
    "        bit = BIT(n)\n",
    "        new = sorted([num - i for i, num in enumerate(nums)])\n",
    "        res = - float('inf')\n",
    "        for i, num in enumerate(nums):\n",
    "            idx = bisect.bisect_left(new, num - i)\n",
    "            cur = max(0, bit.query(idx)) + num \n",
    "            bit.update(idx, cur)\n",
    "            res = max(res, cur)\n",
    "        return res \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class BIT:\n",
    "    def __init__(self, n):\n",
    "        self.array = [- float('inf')] * (n + 1)\n",
    "    \n",
    "    def update(self, idx, val):\n",
    "        idx += 1 \n",
    "        while idx < len(self.array):\n",
    "            self.array[idx] = max(self.array[idx], val)\n",
    "            idx += idx & (- idx)\n",
    "    \n",
    "    def query(self, idx):\n",
    "        res = - float('inf')\n",
    "        idx += 1 \n",
    "        while idx > 0:\n",
    "            res = max(res, self.array[idx])\n",
    "            idx -= idx & (- idx)\n",
    "        return res \n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:\n",
    "        n = len(nums)\n",
    "        new = sorted([num - i for i, num in enumerate(nums)])\n",
    "        bit = BIT(n)\n",
    "        # res = - float('inf')\n",
    "        for i, num in enumerate(nums):\n",
    "            idx = bisect_left(new, num - i)\n",
    "            cur = max(bit.query(idx), 0) + num \n",
    "            # res = max(res, cur)\n",
    "            bit.update(idx, cur)\n",
    "        return bit.query(n - 1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:\n",
    "        if all(num < 0 for num in nums):\n",
    "            return max(nums)\n",
    "        b = sorted([v-i for i, v in enumerate(nums) if v >= 0])\n",
    "        d = {v:i for i, v in enumerate(b)}\n",
    "        t = BIT(len(b))\n",
    "        ans = 0\n",
    "        for i, v in enumerate(nums):\n",
    "            if v < 0: continue\n",
    "            mx = max(0, t.pmax(d[v-i]+1))\n",
    "            t.update(d[v-i]+1, mx + v)\n",
    "            ans = max(ans, mx+v)\n",
    "        return ans\n",
    "\n",
    "\n",
    "class BIT:\n",
    "    def __init__(self, n):\n",
    "        self.n = n\n",
    "        self.tree = [-inf] * (n+1)\n",
    "    \n",
    "    def update(self, x, val):\n",
    "        while x < len(self.tree):\n",
    "            self.tree[x] = max(self.tree[x], val)\n",
    "            x += x & -x\n",
    "    \n",
    "    def pmax(self, x):\n",
    "        ret = -inf\n",
    "        while x > 0:\n",
    "            ret = max(ret, self.tree[x])\n",
    "            x -= x & -x\n",
    "        return ret\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:\n",
    "        s = []\n",
    "\n",
    "        res = nums[0]\n",
    "        \n",
    "        for i, j in enumerate(nums):\n",
    "            x = j - i\n",
    "            \n",
    "            # if not s:\n",
    "            #     s.append([x, j])\n",
    "            \n",
    "            index = bisect_right(s, [x, float('inf')])\n",
    "            \n",
    "            if index == 0:\n",
    "                s.insert(index, [x, j])\n",
    "            else:\n",
    "                s.insert(index, [x, j + max(0, s[index - 1][1])])\n",
    "            \n",
    "            while len(s) > index + 1 and s[index + 1][1] <= s[index][1]:\n",
    "                del s[index + 1]\n",
    "\n",
    "            res = max(res, s[-1][1])\n",
    "            \n",
    "            while s and s[-1][1] < 0:\n",
    "                s.pop(-1)\n",
    "            # print(s)\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "\"\"\"\n",
    "    nums[i] - nums[j] >= i - j  j<i\n",
    "    nums[i] - i >= nums[j] - j\n",
    "    定义 b[i] = nums[i] - i\n",
    "    b[i] >= b[j]    \n",
    "    从 b 中选一个子序列，满足这个子序列是一个非降的序列\n",
    "    求对应的 nums 的元素和的最大值\n",
    "\n",
    "    f[i] = 以下标 i 结尾的子序列，对应的 nums 的元素和的最大值\n",
    "    答案 = max(f)\n",
    "\n",
    "    状态转移方程\n",
    "    f[i] = max(f[j] + nums[i]) for j\n",
    "    \n",
    "\n",
    "    以下标 i 结尾的子序列，对应的 nums 的元素和的最大值 = max(以下标 j 结尾的子序列，对应的 nums 的元素和的最大值 + nums[i]) for j\n",
    "    j 满足： j < i 并且 b[j] <= b[i] (转移来源在值域上是连续的)\n",
    "\n",
    "    下标上不一定是连续的，但是值域上一定是连续的\n",
    "\n",
    "    ------------------------------------------------------------------------\n",
    "    状态转移方程\n",
    "    f[bi] = max(f[bj] + nums[i],0) for bj\n",
    "    j 满足：b[j] <= b[i] (转移来源在值域上是连续的)\n",
    "\n",
    "    假设我们有一个数组，这个数组的下标 p 维护着 b[i] = p 的最大值 f[i]\n",
    "    维护小于等于bi的那些数对应的那个fi就好了\n",
    "\n",
    "\n",
    "    \n",
    "\"\"\"\n",
    "class Solution:\n",
    "    def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:\n",
    "        b = sorted(set(x-i for i,x in enumerate(nums))) # 离散化:复制，排序，去重\n",
    "        t = BIT(len(b)+1)\n",
    "        ans = -inf\n",
    "        for i,x in enumerate(nums):\n",
    "            j = bisect_left(b,x-i)+1 # nums[i]-1 离散化后的值 \n",
    "            fi = max(t.pre_max(j),0)+x\n",
    "            ans = max(ans,fi)\n",
    "            t.update(j,fi)\n",
    "        return ans\n",
    " \n",
    "# 用若干个小区间去合成一个大区间\n",
    "# 树状数组模板（维护前缀最大值）\n",
    "class BIT:\n",
    "    def __init__(self,n:int):\n",
    "        self.tree = [-inf]*n\n",
    "    # low bit去掉最小元素\n",
    "    # 0001110   i\n",
    "    # 1110001\n",
    "    # 1110010  -i\n",
    "\n",
    "    # 0001110    i\n",
    "    # 1110010   -i\n",
    "    # 0000010   i&(-i)\n",
    "    def update(self,i:int,val:int) -> None:\n",
    "        while i < len(self.tree):\n",
    "            self.tree[i] = max(self.tree[i],val);\n",
    "            i += i & -i\n",
    "\n",
    "    def pre_max(self,i:int) -> int:\n",
    "        mx=-inf\n",
    "        while i>0:\n",
    "            mx = max(mx,self.tree[i])\n",
    "            i &=i-1 #等价于 i -= i & -i\n",
    "        return mx\n",
    "\"\"\"\n",
    "    index: 0  1  2  3\n",
    "           3, 3, 5, 6\n",
    "                \n",
    "    nums[2] - nums[0] >= 2 - 0\n",
    "    nums[3] - nums[0] >= 3 - 2\n",
    "    ↓\n",
    "    nums[i] - nums[j] >= i - j\n",
    "\n",
    "\n",
    "    index: 0   1   2   3\n",
    "        5, -1, -3,  8\n",
    "                \n",
    "    nums[3] - nums[0] >= 3 - 0\n",
    "    ↓\n",
    "    nums[i] - nums[j] >= i - j\n",
    "\"\"\"\n",
    "\n",
    "\"\"\"\n",
    "    index:  0  1  2  3\n",
    "            3, 3, 5, 6\n",
    "      b[]:  3, 2, 3, 3   \n",
    "    sum = 3+5+6=14\n",
    "\n",
    "    index: 0   1   2   3\n",
    "           5, -1, -3,  8\n",
    "      b[]: 5, -2, -5,  5   \n",
    "    sum = 5+8=13\n",
    "\n",
    "index   0   1   2   3\n",
    "    a = 2   6   5   7\n",
    "    b = 2   5   3   4\n",
    "\n",
    "假设我们有一个数组，这个数组的下标 p 维护着 b[i] = p 的最大值 f[i]\n",
    "      p: 1     2     3     4     5\n",
    "        -∞     0     1     2     1\n",
    "\n",
    "f3 ← (f0,f2)\n",
    "\n",
    "\n",
    "----------------------------\n",
    "知识点：离散化\n",
    "-3     6     -1     2     2\n",
    "例如映射到 \n",
    "1      4      2     3     3 \n",
    "----------------------------\n",
    "\n",
    "\n",
    "-3     6     -1     2     2\n",
    "↓（排序）\n",
    "-3     -1     2     2     6\n",
    "↓（去重）\n",
    "-3     -1     2     6\n",
    "要实现映射到 \n",
    "1       2     3     4      \n",
    "\n",
    "-3     -1     2     6\n",
    "0       1     2     3   +1\n",
    "1       2     3     4\n",
    "\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "inf = 10**30\n",
    "\n",
    "\n",
    "class Solution:\n",
    "\n",
    "    def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:\n",
    "        b = sorted(set(x - i for i, x in enumerate(nums)))  # 离散化 nums[i]-i\n",
    "        t = BIT(len(b) + 1)\n",
    "        ans = -inf\n",
    "        for i, x in enumerate(nums):\n",
    "            j = bisect_left(b, x - i) + 1  # nums[i]-i 离散化后的值（从 1 开始）\n",
    "            f = max(t.pre_max(j), 0) + x\n",
    "            ans = max(ans, f)\n",
    "            t.update(j, f)\n",
    "        return ans\n",
    "\n",
    "\n",
    "# 树状数组模板（维护前缀最大值）\n",
    "class BIT:\n",
    "\n",
    "    def __init__(self, n: int):\n",
    "        self.tree = [-inf] * n\n",
    "\n",
    "    def update(self, i: int, val: int) -> None:\n",
    "        while i < len(self.tree):\n",
    "            self.tree[i] = max(self.tree[i], val)\n",
    "            i += i & -i\n",
    "\n",
    "    def pre_max(self, i: int) -> int:\n",
    "        mx = -inf\n",
    "        while i > 0:\n",
    "            mx = max(mx, self.tree[i])\n",
    "            i &= i - 1\n",
    "        return mx\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "\"\"\"\n",
    "由题意可以得到nums[i]-i >= nums[j]-j\n",
    "将nums[i]处理为nums[i]-i\n",
    "这样的话子序列必须是非递减的\n",
    "先处理更小的更靠前的数\n",
    "设dp[i]为以第i个元素结尾的子序列的最大和\n",
    "\"\"\"\n",
    "class BIT:\n",
    "    def __init__(self, n: int):\n",
    "        self.tree = [-inf] * n\n",
    "\n",
    "    def update(self, i: int, val: int) -> None:\n",
    "        while i < len(self.tree):\n",
    "            self.tree[i] = max(self.tree[i], val)\n",
    "            i += i & -i\n",
    "\n",
    "    def pre_max(self, i: int) -> int:\n",
    "        mx = -inf\n",
    "        while i > 0:\n",
    "            mx = max(mx, self.tree[i])\n",
    "            i &= i - 1\n",
    "        return mx\n",
    "\n",
    "class Solution:\n",
    "    def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:\n",
    "        n = len(nums)\n",
    "        bit = BIT(n)\n",
    "        for i, num in enumerate(nums):\n",
    "            nums[i] = num-i\n",
    "        sorted_idxs = sorted(range(n), key=lambda i:(nums[i], i))\n",
    "        dp = [0] * n\n",
    "        for idx in sorted_idxs:\n",
    "            if idx == 0:\n",
    "                dp[0] = nums[0]\n",
    "                bit.update(1, dp[0])\n",
    "                continue\n",
    "            mx = bit.pre_max(idx)\n",
    "            dp[idx] = max(0, mx)+nums[idx]+idx\n",
    "            bit.update(idx+1, dp[idx])\n",
    "        return max(dp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from sortedcontainers import SortedList\n",
    "\n",
    "class Solution:\n",
    "    def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:\n",
    "        dp = SortedList()\n",
    "        map = {}\n",
    "        for i, x in enumerate(nums):\n",
    "            p = bisect_left(dp, x - i)\n",
    "            if p < len(dp) and dp[p] == x - i:\n",
    "                if map[dp[p]] < 0:\n",
    "                    map[dp[p]] = max(map[dp[p]], x)\n",
    "                elif x > 0:\n",
    "                    map[dp[p]] += x                 \n",
    "            else:\n",
    "                dp.add(x - i)\n",
    "                if p > 0:\n",
    "                    map[x-i] = x + max(map[dp[p - 1]], 0)\n",
    "                else:\n",
    "                    map[x-i] = x\n",
    "            c = map[dp[p]]\n",
    "            while p + 1 < len(dp) and map[dp[p + 1]] < c:\n",
    "                dp.pop(index=p+1)   \n",
    "\n",
    "        return map[dp[-1]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class SegmentTree:\n",
    "    def __init__(self, nums: List[int], merge=max):\n",
    "        self.nums = nums\n",
    "        self.n = len(nums)\n",
    "        self.merge = merge\n",
    "        self.tree = [None] * (self.n * 4)\n",
    "        if self.n:\n",
    "            self.build(0, 0, self.n - 1)\n",
    "    def build(self, o: int, l: int, r: int):\n",
    "        if l == r:\n",
    "            self.tree[o] = self.nums[l]\n",
    "            return\n",
    "        mid = (l + r) // 2\n",
    "        lo, ro = o * 2 + 1, o * 2 + 2\n",
    "        self.build(lo, l, mid)\n",
    "        self.build(ro, mid + 1, r)\n",
    "        self.tree[o] = self.merge(self.tree[lo], self.tree[ro])\n",
    "\n",
    "    def update(self, idx: int, val: int):\n",
    "        self.nums[idx] = val\n",
    "        self._update(0, 0, self.n - 1, idx)\n",
    "    def _update(self, o: int, l: int, r: int, idx: int):\n",
    "        if l == r == idx:\n",
    "            self.tree[o] = self.nums[idx]\n",
    "            return\n",
    "        lo, ro = o * 2 + 1, o * 2 + 2\n",
    "        mid = (l + r) // 2\n",
    "        if idx <= mid:\n",
    "            self._update(lo, l, mid, idx)\n",
    "        else:\n",
    "            self._update(ro, mid + 1, r, idx)\n",
    "        self.tree[o] = self.merge(self.tree[lo], self.tree[ro])\n",
    "\n",
    "    def query(self, L: int, R: int) -> int:\n",
    "        return self._query(0, 0, self.n - 1, L, R)\n",
    "    def _query(self,o: int, l: int, r: int, L: int, R: int) -> int:\n",
    "        if L <= l and r <= R:\n",
    "            return self.tree[o]\n",
    "        lo, ro = o * 2 + 1, o * 2 + 2\n",
    "        mid = (l + r) // 2\n",
    "        if mid >= R:\n",
    "            return self._query(lo, l, mid, L, R)\n",
    "        if mid + 1 <= L:\n",
    "            return self._query(ro, mid + 1, r, L, R)\n",
    "        return self.merge(self._query(lo, l, mid, L, R), self._query(ro, mid + 1, r, L, R))\n",
    "\n",
    "        \n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:\n",
    "        ans = -inf\n",
    "        n = len(nums)\n",
    "        for i in range(n):\n",
    "            nums[i] -= i\n",
    "        # 归一化\n",
    "        a = sorted(nums, key=lambda x: x)\n",
    "        cnt = {}\n",
    "        for i, x in enumerate(a):\n",
    "            # 相同数字归一到最大的下标索引\n",
    "            cnt[x] = i\n",
    "        tree = SegmentTree([-inf] * len(a), max)\n",
    "        #print(f'nums:{nums}')\n",
    "        #print(f'cnt:{cnt}')\n",
    "        #print(f'a:{a}')\n",
    "        for i, x in enumerate(nums):\n",
    "            premx = tree.query(0, cnt[x])\n",
    "            sm = x + i  + max(0, premx)\n",
    "            tree.update(cnt[x], sm)\n",
    "            #print(f'ans:{ans},sm:{sm},premx:{premx},i:{i},x:{x},cnt[x]:{cnt[x]}')\n",
    "            ans = max(ans, sm)\n",
    "        #print(f'ans:{ans}')\n",
    "        return ans\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from sortedcontainers import SortedList\n",
    "class Solution:\n",
    "    def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:\n",
    "        tmp = all([nu <= 0 for nu in nums])\n",
    "        if tmp:\n",
    "            return max(nums)\n",
    "        delta = [(x - i, x) for i, x in enumerate(nums) if x > 0]\n",
    "        # print(delta)\n",
    "        sl = SortedList()\n",
    "        for x, y in delta:\n",
    "            tmp = sl.bisect_right((x, inf))\n",
    "            val = (sl[tmp - 1][1] if tmp else 0) + y\n",
    "            while tmp < len(sl):\n",
    "                if sl[tmp][1] < val:\n",
    "                    sl.remove(sl[tmp])\n",
    "                else:\n",
    "                    break\n",
    "            sl.add((x, val))\n",
    "        return sl[-1][1]\n",
    "                \n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class SegTree:\n",
    "\n",
    "    __slots__ = {'n', 'op', 'e', 'log', 'size', 'd'}\n",
    "\n",
    "    def __init__(self, V, OP, E):\n",
    "        '''\n",
    "        V: 原始数组\n",
    "        OP: 维护的运算(min, max, sum...)\n",
    "        E: 线段树初值\n",
    "        '''\n",
    "        self.n = len(V)\n",
    "        self.op = OP\n",
    "        self.e = E\n",
    "        self.log = (self.n - 1).bit_length()\n",
    "        self.size = 1 << self.log\n",
    "        self.d=[E for i in range(2*self.size)]\n",
    "        # for i in range(self.n):\n",
    "        #     self.d[self.size + i] = V[i]\n",
    "        # for i in range(self.size - 1, 0, -1):\n",
    "        #     self.update(i)\n",
    "\n",
    "    def set(self, p, x):\n",
    "        assert 0 <= p and p < self.n\n",
    "        p += self.size\n",
    "        self.d[p] = x\n",
    "        for i in range(1, self.log + 1):\n",
    "            self.update(p >> i)\n",
    "\n",
    "    def get(self, p):\n",
    "        assert 0 <= p and p <self.n\n",
    "        return self.d[p + self.size]\n",
    "\n",
    "    def prod(self, l ,r):\n",
    "        #[l, r)\n",
    "        assert 0 <= l and l <= r and r <= self.n\n",
    "        sml = self.e\n",
    "        smr = self.e\n",
    "        l += self.size\n",
    "        r += self.size\n",
    "\n",
    "        while(l < r):\n",
    "            if (l & 1):\n",
    "                sml = self.op(sml, self.d[l])\n",
    "                l += 1\n",
    "            if (r & 1):\n",
    "                smr = self.op(self.d[r - 1], smr)\n",
    "                r -= 1\n",
    "            l >>= 1\n",
    "            r >>= 1\n",
    "        return self.op(sml, smr)\n",
    "\n",
    "    def all_prod(self):\n",
    "        return self.d[1]\n",
    "\n",
    "    def max_right(self, l, f):\n",
    "        assert 0 <= l and l <= self.n\n",
    "        assert f(self.e)\n",
    "        if l == self.n:\n",
    "            return self.n\n",
    "        l += self.size\n",
    "        sm = self.e\n",
    "        while(1):\n",
    "            while(l % 2 == 0):\n",
    "                l >>= 1\n",
    "            if not(f(self.op(sm, self.d[l]))):\n",
    "                while(l < self.size):\n",
    "                    l = 2 * l\n",
    "                    if f(self.op(sm, self.d[l])):\n",
    "                        sm = self.op(sm, self.d[l])\n",
    "                        l += 1\n",
    "                return l - self.size\n",
    "            sm = self.op(sm,self.d[l])\n",
    "            l += 1\n",
    "            if (l & -l) == l:\n",
    "                break\n",
    "        return self.n\n",
    "    def min_left(self, r, f):\n",
    "        assert 0 <= r and r < self.n\n",
    "        assert f(self.e)\n",
    "        if r == 0:\n",
    "            return 0\n",
    "        r += self.size\n",
    "        sm = self.e\n",
    "        while(1):\n",
    "            r -= 1\n",
    "            while(r > 1 & (r % 2)):\n",
    "                r >>= 1\n",
    "            if not(f(self.op(self.d[r], sm))):\n",
    "                while(r < self.size):\n",
    "                    r = (2 * r + 1)\n",
    "                    if f(self.op(self.d[r], sm)):\n",
    "                        sm=self.op(self.d[r], sm)\n",
    "                        r -= 1\n",
    "                return r + 1 - self.size\n",
    "            sm = self.op(self.d[r] ,sm)\n",
    "            if (r & -r) == r:\n",
    "                break\n",
    "        return 0\n",
    "\n",
    "    def update(self, k):\n",
    "        self.d[k] = self.op(self.d[2 * k], self.d[2 * k + 1])\n",
    "        \n",
    "    def __str__(self):\n",
    "        return str([self.get(i) for i in range(self.n)])\n",
    "\n",
    "class Solution:\n",
    "    def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:\n",
    "        if all(num < 0 for num in nums):\n",
    "            return max(nums)\n",
    "        n = len(nums)\n",
    "        tmp = [v - i for i, v in enumerate(nums)]\n",
    "        d = {v : i + 1 for i, v in enumerate (list(sorted(set(tmp))))}\n",
    "\n",
    "        sg = SegTree([0 for _ in range(len(d) + 10)], lambda x, y: x if x > y else y, 0)\n",
    "\n",
    "        for i in range(n):\n",
    "            idx = d[tmp[i]]\n",
    "            cur = sg.prod(1, idx + 1)\n",
    "            sg.set(idx, cur + nums[i] if nums[i] > 0 else 0)\n",
    "        \n",
    "        return sg.all_prod()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:\n",
    "        n = len(nums)\n",
    "        print([nums[i] - i for i in range(n)])\n",
    "        b = sorted(list(set([nums[i] - i for i in range(n)])))\n",
    "        d = dict()\n",
    "        for i in range(len(b)):\n",
    "            d[b[i]] = i + 1\n",
    "        bit_tree = BIT(n + 1)\n",
    "        ans = -inf\n",
    "        for i, x in enumerate(nums):\n",
    "            idx = d[x-i]\n",
    "            f = max(0, bit_tree.pre_max(idx)) + x\n",
    "            ans = max(ans, f)\n",
    "            bit_tree.update(idx, f)\n",
    "        return ans\n",
    "\n",
    "\n",
    "\n",
    "class BIT:\n",
    "\n",
    "    def __init__(self, n: int):\n",
    "        self.tree = [-inf] * n\n",
    "\n",
    "    def pre_max(self, i):\n",
    "        mx = -inf\n",
    "        while i > 0:\n",
    "            mx = max(mx, self.tree[i])\n",
    "            i -= i & -i\n",
    "        return mx\n",
    "    \n",
    "    def update(self, i, val):\n",
    "        while i < len(self.tree):\n",
    "            self.tree[i] = max(self.tree[i], val)\n",
    "            i += i & -i\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:\n",
    "        inf = 10 ** 9 + 1\n",
    "        n = len(nums)\n",
    "        arr = [0] * n\n",
    "        #  离散化\n",
    "        for i in range(n):\n",
    "            arr[i] = nums[i] - i\n",
    "        arr = sorted(set(arr))\n",
    "        idxs = {}\n",
    "        i = 1\n",
    "        for a in arr:\n",
    "            idxs[a] = i\n",
    "            i += 1\n",
    "\n",
    "        #  树状模板\n",
    "        n = len(arr)\n",
    "        tree = [-inf] * (n + 1)\n",
    "        def update(i, val):\n",
    "            while i <= n:\n",
    "                tree[i] = max(tree[i], val)\n",
    "                i += i & -i\n",
    "        def query(i):\n",
    "            mx = -inf\n",
    "            while i > 0:\n",
    "                mx = max(tree[i], mx)\n",
    "                i -= i & -i \n",
    "            return mx\n",
    "\n",
    "        #  逐个查询\n",
    "        for i, x in enumerate(nums):\n",
    "            j = idxs[x - i]\n",
    "            f = max(query(j), 0) + x\n",
    "            update(j, f)\n",
    "        \n",
    "        return query(n)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class LineTree:\n",
    "    def __init__(self, nums):\n",
    "        self.n = len(nums)\n",
    "        self.tree = [0] * (4 * self.n + 1)\n",
    "        self.build(nums, 1, 0, self.n-1)\n",
    "    \n",
    "    def build(self, nums, pos, l, r):\n",
    "        if l == r:\n",
    "            self.tree[pos] = nums[l]\n",
    "            return\n",
    "        mid = (l + r) >> 1\n",
    "        l_child = 2 * pos\n",
    "        r_child = 2 * pos + 1\n",
    "        self.build(nums, l_child, l, mid)\n",
    "        self.build(nums, r_child, mid+1, r)\n",
    "        self.tree[pos] = max(self.tree[l_child], self.tree[r_child])\n",
    "\n",
    "    def update(self, pos, l, r, k, val):\n",
    "        if l == r:\n",
    "            self.tree[pos] = val\n",
    "            return\n",
    "        mid = (l + r) >> 1\n",
    "        if k <= mid:\n",
    "            self.update(pos * 2, l, mid, k, val)\n",
    "        else:\n",
    "            self.update(pos * 2 + 1, mid + 1, r, k, val)\n",
    "        self.tree[pos] = max(self.tree[pos * 2], self.tree[pos * 2 + 1])\n",
    "\n",
    "    def find(self, pos, l, r, t_l, t_r):\n",
    "        if t_l <= l and r <= t_r:\n",
    "            return self.tree[pos]\n",
    "        mid = (l + r) >> 1\n",
    "        if t_r <= mid:\n",
    "            return self.find(pos * 2, l, mid, t_l, t_r)\n",
    "        elif t_l > mid:\n",
    "            return self.find(pos * 2 + 1, mid + 1, r, t_l, t_r)\n",
    "        return max(self.find(pos * 2, l, mid, t_l, t_r), self.find(pos * 2 + 1, mid + 1, r, t_l, t_r))\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:\n",
    "        all_diff = set([nums[i] - i for i in range(len(nums))])\n",
    "        all_diff = sorted(list(all_diff))\n",
    "        diff_map = {k: v for v, k in enumerate(all_diff)}\n",
    "        dp = [int(-1e9-1)] * len(diff_map)\n",
    "        seg_tree = LineTree(dp)\n",
    "        for i, n in enumerate(nums):\n",
    "            idx = diff_map[n - i]\n",
    "            prev_max = seg_tree.find(1, 0, len(dp) - 1, 0, idx)\n",
    "            curr_max = max(prev_max, 0) + n\n",
    "            seg_tree.update(1, 0, len(dp) - 1, idx, curr_max)\n",
    "        return seg_tree.tree[1]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class BIT: # 维护前缀最大值\n",
    "    def __init__(self, n):\n",
    "        self.n = n\n",
    "        self.c = [-inf for _ in range(n + 1)] # 树状数组从下标1开始存储，表示原数组a下标区间(i-lowbit(i), i]的max（原数组a也从下标1开始存储）\n",
    "\n",
    "    def _lowbit(self, x): # lowbit指x二进制下最低位的1和后面所有的0组成的数\n",
    "        return x & (-x)\n",
    "\n",
    "    # 将原数组a[i]的值变为x(如果)，修改树状数组对应的位置\n",
    "    def update(self, i, x):\n",
    "        while i <= self.n:\n",
    "            self.c[i] = max(self.c[i], x)\n",
    "            i += self._lowbit(i)\n",
    "\n",
    "    # 求原数组a下标区间[1,i]的max（前i项的max）\n",
    "    def pre_max(self, i):\n",
    "        mx = -inf\n",
    "        while i > 0:\n",
    "            mx = max(mx, self.c[i])\n",
    "            i -= self._lowbit(i)\n",
    "        return mx\n",
    "\n",
    "class Solution:\n",
    "    def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:\n",
    "        b = sorted(set(x - i for i, x in enumerate(nums)))  # 离散化 nums[i]-i\n",
    "        d = {x: i + 1 for i, x in enumerate(b)}\n",
    "\n",
    "        t = BIT(len(b))\n",
    "        for i, x in enumerate(nums):\n",
    "            j = d[x - i]  # nums[i]-i 离散化后的值（从 1 开始）\n",
    "            f = max(t.pre_max(j), 0) + x\n",
    "            t.update(j, f)\n",
    "        return t.pre_max(len(b))  # 所有 f 的最大值\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:\n",
    "        \"\"\"\n",
    "        假设线段树要维护的数组为f，f的区间下标是离散化之后（由小到大映射到1,2,3...）的x - i，f[c]表示以x结尾，x - i = c的子序列，对应的nums元素和的最大值。\n",
    "        我们要维护一个前缀最大值数组premax，premax（L， R）表示[L, R]区间内f数组元素的最大值\n",
    "        有状态转移方程：\n",
    "            f[c] = max(premax(1, c - 1), 0) + x\n",
    "        \"\"\"\n",
    "        b = sorted(set(x - i for i,x in enumerate(nums)))\n",
    "        n = len(b)\n",
    "        premax = [-inf] * (n * 4)\n",
    "\n",
    "        def update(o, l, r, L, R, val):\n",
    "            if L <= l and r <= R:\n",
    "                premax[o] = max(premax[o], val)\n",
    "                return \n",
    "            m = (l + r) // 2\n",
    "            if m >= L: update(o * 2, l, m, L, R, val)\n",
    "            if m < R: update(o * 2 + 1, m + 1, r, L, R, val)\n",
    "            premax[o] = max(premax[o * 2], premax[o * 2 + 1])\n",
    "            \n",
    "        \n",
    "        # 查询[L,R]区间f数组元素的最大值\n",
    "        def query(o, l, r, L, R):\n",
    "            if L > R:\n",
    "                return -inf\n",
    "            if L <= l and r <= R:\n",
    "                return premax[o]\n",
    "            res = -inf\n",
    "            m = (l + r) // 2\n",
    "            if m >= L: res = max(res, query(o * 2, l, m, L, R))\n",
    "            if m < R: res = max(res, query(o * 2 + 1, m + 1, r, L, R))\n",
    "            return res\n",
    "\n",
    "        ans = -inf\n",
    "        for i, x in enumerate(nums):\n",
    "            j = bisect_left(b, x - i) + 1\n",
    "            f = max(query(1, 1, n, 1, j), 0) + x\n",
    "            ans = max(ans, f)\n",
    "            update(1, 1, n, j, j, f)\n",
    "\n",
    "        return ans\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class SegmentTree:\n",
    "    def __init__(self, arr):\n",
    "        n = len(arr)\n",
    "        self.arr = [0] + arr\n",
    "        k = math.ceil(math.log2(n))\n",
    "        self.max = [0] * (1 << k + 1)\n",
    "        self.lazy = [0] * (1 << k + 1)\n",
    "\n",
    "    def build(self, k, l, r):\n",
    "        if l == r:\n",
    "            self.max[k] = self.arr[l]\n",
    "        else:\n",
    "            m = l + (r - l) // 2\n",
    "            self.build(k << 1, l, m)\n",
    "            self.build(k << 1 | 1, m + 1, r)\n",
    "            self.pushup(k)\n",
    "\n",
    "    # define pushup function based on problems\n",
    "    def pushup(self, k):\n",
    "        self.max[k] = max(self.max[k << 1], self.max[k << 1 | 1])\n",
    "\n",
    "    def pushdown2(self, k, l, r, m):\n",
    "        if self.lazy[k]:\n",
    "            self.lazy[k << 1] = self.lazy[k << 1 | 1] = self.lazy[k]\n",
    "            self.max[k << 1] = self.lazy[k]\n",
    "            self.max[k << 1 | 1] = self.lazy[k]\n",
    "            self.lazy[k] = 0\n",
    "\n",
    "    def update(self, L, R, v, l, r, k):\n",
    "        if L <= l and r <= R:\n",
    "            self.lazy[k] = v\n",
    "            self.max[k] = v\n",
    "        else:\n",
    "            m = l + (r - l) // 2\n",
    "            self.pushdown2(k, l, r, m)\n",
    "            if L <= m:\n",
    "                self.update(L, R, v, l, m, k << 1)\n",
    "            if R > m:\n",
    "                self.update(L, R, v, m + 1, r, k << 1 | 1)\n",
    "            self.pushup(k)\n",
    "\n",
    "    # define query function based on problems\n",
    "    def queryMax(self, L, R, l, r, k):\n",
    "        if L <= l and r <= R:\n",
    "            return self.max[k]\n",
    "        else:\n",
    "            m = l + (r - l) // 2\n",
    "            self.pushdown2(k, l, r, m) # add: pushdown; update: pushdown2\n",
    "            ans = float(\"-inf\")\n",
    "            if L <= m:\n",
    "                ans = max(ans, self.queryMax(L, R, l, m, k << 1))\n",
    "            if R > m:\n",
    "                ans = max(ans, self.queryMax(L, R, m + 1, r, k << 1 | 1))\n",
    "            return ans\n",
    "\n",
    "from sortedcontainers import SortedList\n",
    "class Solution:\n",
    "    def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:\n",
    "        n = len(nums)\n",
    "        A = SortedList()\n",
    "        dic = defaultdict(lambda: -inf)\n",
    "        ans = -inf\n",
    "        B = sorted(set([a - i for i, a in enumerate(nums)]))\n",
    "        print(B)\n",
    "        dic2 = {b: i for i, b in enumerate(B)}\n",
    "        m = len(dic2)\n",
    "        C = [-inf] * m\n",
    "        sg = SegmentTree(C)\n",
    "        sg.build(1, 1, m)\n",
    "        for i, a in enumerate(nums):\n",
    "            x = a - i\n",
    "            idx = A.bisect_right(x)\n",
    "            if idx:\n",
    "                y = A[idx - 1]\n",
    "                y = dic2[y]\n",
    "                z = sg.queryMax(1, y + 1, 1, m, 1)\n",
    "                tmp = max(0, z) + a\n",
    "            else:\n",
    "                tmp = a\n",
    "            j = dic2[x]\n",
    "            sg.update(j + 1, j + 1, tmp, 1, m, 1)\n",
    "\n",
    "            ans = max(ans, tmp)\n",
    "            A.add(x)\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class BIT:\n",
    "    def __init__(self, nn):\n",
    "        self.n = nn + 1\n",
    "        self.tree = [0] * self.n\n",
    "    def update(self, i: int, x: int):\n",
    "        while i < self.n:\n",
    "            self.tree[i] = max(self.tree[i], x)\n",
    "            i += i & -i\n",
    "    def query(self, i: int):\n",
    "        res = 0\n",
    "        while i:\n",
    "            res = max(res, self.tree[i])\n",
    "            i -= i & -i\n",
    "        return res\n",
    "class Solution:\n",
    "    def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:\n",
    "        n = len(nums)\n",
    "        f = [[nums[i] - i, nums[i]]for i in range(n)]\n",
    "        \"\"\"\n",
    "        (3,3) (2,3) (3,5) (3,6)\n",
    "        (4,3) (1,3) (4,3) (4,3)\n",
    "\n",
    "        1     4\n",
    "        0     0  (4,3)\n",
    "        0     3  (1,3)\n",
    "        3     3  (4,6)\n",
    "        3     6  (4, 9)\n",
    "\n",
    "        \"\"\"\n",
    "        a = sorted(f, key = lambda x: x[0])\n",
    "        mp = defaultdict(int)\n",
    "        for i in range(n):\n",
    "            mp[a[i][0]] = i + 1\n",
    "        for i in range(n):\n",
    "            f[i][0] = mp[f[i][0]]\n",
    "        print(f)\n",
    "        ans = -inf\n",
    "        root = BIT(n)\n",
    "        for i in range(n):\n",
    "            x = f[i][0]\n",
    "            val = f[i][1]\n",
    "            res = max(val, root.query(x) + val)\n",
    "            print(x, res)\n",
    "            root.update(x, res)\n",
    "            ans = max(ans, res)\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "inf = 10 ** 14 + 1\n",
    "class BIT:\n",
    "    def __init__(self, n):\n",
    "        self.n = n\n",
    "        self.tr = [-inf]*(n + 1)\n",
    "        self.a = [-inf] * (n + 1)\n",
    "\n",
    "    def update(self, x, v):\n",
    "        lx = x\n",
    "        if v > self.a[x]:\n",
    "            self.a[x] = v\n",
    "            while x <= self.n:                \n",
    "                self.tr[x] = self.a[lx] if self.a[lx] > self.tr[x] else self.tr[x]\n",
    "                x += -x&x\n",
    "\n",
    "    def query(self, x, y):\n",
    "        ans = -inf\n",
    "        \n",
    "        while x <= y:\n",
    "            ans = self.a[y] if self.a[y] > ans else ans\n",
    "            y -= 1\n",
    "            while y - (-y&y) >= x:\n",
    "                ans = self.tr[y] if self.tr[y] > ans else ans\n",
    "                y -= -y&y\n",
    "        return ans\n",
    "class Solution:\n",
    "    def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:\n",
    "        unique = set()\n",
    "        n = len(nums)\n",
    "        for i in range(1, n + 1):\n",
    "            unique.add(nums[i-1] - i)\n",
    "            # unique.add(nums[i-1])\n",
    "        \n",
    "        unique = list(sorted(unique))\n",
    "        # print(unique)\n",
    "        num2ind = {}\n",
    "        ind2num = {}\n",
    "        for ind, val in enumerate(unique, 1):\n",
    "            num2ind[val] = ind\n",
    "            ind2num[ind] = val\n",
    "        N = ind\n",
    "        bit = BIT(N)\n",
    "        # nums1 = [num2ind[nums[i] - i - 1] for i in range(n-1, -1, -1)]\n",
    "        # print(nums1)\n",
    "        for i in range(n):\n",
    "            val1 = nums[i] - i - 1\n",
    "            # print(nums[i], i)\n",
    "            val2 = nums[i]\n",
    "            ind1 = num2ind[val1]\n",
    "            val = bit.query(1, ind1)\n",
    "            # print(val, val2, nums[i], ind1)\n",
    "            bit.update(ind1, max(val, 0) + val2)\n",
    "        \n",
    "        return bit.query(1, N)\n",
    "            \n",
    "        \n",
    "            \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:\n",
    "        \n",
    "        m = {}\n",
    "        n = 1\n",
    "        for b in sorted(v - i for i, v in enumerate(nums)):\n",
    "            if b in m:\n",
    "                continue\n",
    "            \n",
    "            m[b] = n\n",
    "            n += 1\n",
    "\n",
    "       \n",
    "        s = [-1000_000_001] * (2 << n.bit_length())\n",
    "        def add(p, l, r, t, val):\n",
    "            if l == r:\n",
    "                s[p] = val\n",
    "                return\n",
    "            m = (l + r) // 2\n",
    "            if t <= m:\n",
    "                add(p<<1, l, m, t, val)\n",
    "            else:\n",
    "                add(p<<1|1, m+1, r, t, val)\n",
    "            s[p] = max(s[p<<1], s[p<<1|1])\n",
    "\n",
    "        def query(p, l, r, L, R):\n",
    "            if L <= l and r <= R:\n",
    "                return s[p]\n",
    "            \n",
    "            res = -1000_000_001\n",
    "            m = (l + r) // 2\n",
    "\n",
    "            if L <= m:\n",
    "                res = max(res, query(p<<1, l, m, L, R))\n",
    "            \n",
    "            if R > m:\n",
    "                res = max(res, query(p<<1|1, m+1, r, L, R))\n",
    "            \n",
    "            return res\n",
    "\n",
    "        # print(f'{n=} {m=}')\n",
    "        ans = -1000_000_001\n",
    "        for i, v in enumerate(nums):\n",
    "            q = query(1, 1, n, 1, m[v-i])\n",
    "            if q == -1000_000_001:\n",
    "                q = 0\n",
    "            # print(f'sum[{1}~{n}] => {q=} v+q={v+q} v-i={v-i}')\n",
    "            cmp = max(v + q, v)\n",
    "            add(1, 1, n, m[v-i], cmp)\n",
    "            ans = max(ans, cmp)\n",
    "\n",
    "        return ans\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:\n",
    "        arr=sorted(list(set(a-i for i,a in enumerate(nums))))\n",
    "        mp={}\n",
    "        for i,a in enumerate(arr):\n",
    "            mp[a]=i\n",
    "        nums=[[mp[a-i],a] for i,a in enumerate(nums)]\n",
    "        arr=[-inf]*(len(arr)+1)\n",
    "        def lowbit(x):\n",
    "            return x&-x\n",
    "        def getMax(i):\n",
    "            res=0\n",
    "            while i:\n",
    "                res=max(res,arr[i])\n",
    "                i-=lowbit(i)\n",
    "            return max(res,0)\n",
    "        def update(i,x):\n",
    "            while i<len(arr):\n",
    "                arr[i]=max(arr[i],x)\n",
    "                i+=lowbit(i)\n",
    "        for i in range(len(nums)):\n",
    "            nums[i][1]+=getMax(nums[i][0]+1)\n",
    "            update(nums[i][0]+1,nums[i][1])\n",
    "        return max(a[1] for a in nums)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from sortedcontainers import SortedList\n",
    "class Solution:\n",
    "    def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:\n",
    "        L=set()\n",
    "        for i in range(len(nums)):\n",
    "            L.add((nums[i]-i))\n",
    "        L=list(L)\n",
    "        L.sort()\n",
    "        dic=dict()\n",
    "        for i in range(len(L)):\n",
    "            dic[L[i]]=i\n",
    "        lenth=len(L)\n",
    "        tree=[0]*(4*lenth)\n",
    "        def update(i,val,n,left,right):\n",
    "            if right-left==1 and left==i:\n",
    "                tree[n]=max(val,tree[n])\n",
    "                return tree[n]\n",
    "            mid=(left+right)//2\n",
    "            if mid<=i:\n",
    "                son=update(i,val,2*n+2,mid,right)\n",
    "            else:\n",
    "                son=update(i,val,2*n+1,left,mid) \n",
    "            tree[n]=max(son,tree[n])\n",
    "            return tree[n]\n",
    "        def find(n,l,r,tar):\n",
    "            if r==tar:\n",
    "                return tree[n]\n",
    "            m=(r+l)//2\n",
    "            if tar>m:\n",
    "                return max(find(2*n+1,l,m,m),find(2*n+2,m,r,tar))\n",
    "            else:\n",
    "                return find(2*n+1,l,m,tar)\n",
    "        maximum=-10**9\n",
    "        for i in range(len(nums)):\n",
    "            tar=dic[nums[i]-i]\n",
    "            now=nums[i]+find(0,0,lenth,tar+1)\n",
    "            update(tar,now,0,0,lenth)\n",
    "            maximum=max(maximum,now)\n",
    "        return maximum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:\n",
    "        if all(x < 0 for x in nums):\n",
    "            return max(nums)\n",
    "        n = len(nums)\n",
    "        res = []\n",
    "        for i, j in enumerate(nums):\n",
    "            res.append([j - i, i])\n",
    "        res.sort()\n",
    "        que = dict()\n",
    "        for i, (j, k) in enumerate(res):\n",
    "            que[k] = i + 1\n",
    "        mx = n + 1\n",
    "        s = [0] * (4 * mx)\n",
    "\n",
    "        def modify(o, l, r, idx, val):\n",
    "            if l == r:\n",
    "                s[o] = val \n",
    "                return \n",
    "            mid = (l + r) // 2\n",
    "            if idx <= mid:\n",
    "                modify(2 * o, l, mid, idx, val)\n",
    "            else:\n",
    "                modify(2 * o + 1, mid + 1, r, idx, val)\n",
    "            s[o] = max(s[2 * o], s[2 * o + 1])\n",
    "\n",
    "        def get_max(o, l, r, L, R):\n",
    "            if L <= l and r <= R:\n",
    "                return s[o]\n",
    "            mid = (l + r) // 2\n",
    "            res = -inf \n",
    "            if L <= mid:\n",
    "                res = get_max(2 * o, l, mid, L, R)\n",
    "            if R >= mid + 1:\n",
    "                res = max(res, get_max(2 * o + 1, mid + 1, r, L, R))\n",
    "            return res \n",
    "\n",
    "        for i, j in enumerate(nums):\n",
    "            idx = que[i]\n",
    "            pre_max = max(get_max(1, 1, mx, 1, idx), 0)\n",
    "            if pre_max + j > get_max(1, 1, mx, idx, idx):\n",
    "                modify(1, 1, mx, idx, pre_max + j)\n",
    "        return s[1]\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class ST:\n",
    "    def __init__(self):\n",
    "        self.t = {}\n",
    "        \n",
    "    def update(self, l, r, val, s, t, p):\n",
    "        if p not in self.t:\n",
    "            self.t[p] = -inf\n",
    "        if l>t or r<s: return -inf\n",
    "        if l<=s and t<=r:\n",
    "            self.t[p] = val\n",
    "            return\n",
    "        m = (s+t)>>1\n",
    "        self.update(l, r, val, s, m, p<<1)\n",
    "        self.update(l, r, val, m+1, t, p<<1|1)\n",
    "        self.t[p] = max(self.t[p<<1], self.t[p<<1|1])\n",
    "    \n",
    "    def query(self, l, r, s, t, p):\n",
    "        if p not in self.t:\n",
    "            self.t[p] = -inf\n",
    "        if l<=s and t<=r:\n",
    "            return self.t[p]\n",
    "        m = (s+t)>>1\n",
    "        ret = -inf\n",
    "        if l<=m:\n",
    "            ret = self.query(l, r, s, m, p<<1)\n",
    "        if m+1<=r:\n",
    "            ret = max(ret, self.query(l, r, m+1, t, p<<1|1))\n",
    "        return ret\n",
    "\n",
    "class Solution:\n",
    "    def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:\n",
    "        # nums[i]-i >= nums[j]-j\n",
    "        p = set([x-i for i, x in enumerate(nums)])\n",
    "        p = sorted(p)\n",
    "        st = ST()\n",
    "        s = 0\n",
    "        t = len(p)\n",
    "        for i, x in enumerate(nums):\n",
    "            idx = bisect.bisect_right(p, x-i)-1\n",
    "            myv = st.query(0, idx, s, t, 1)\n",
    "            st.update(idx, idx, max(x, myv, x+myv), s, t, 1)\n",
    "        # print(s, t)\n",
    "        return st.query(s, t, s, t, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:\n",
    "        s = [nums[i] - i for i in range(len(nums))]\n",
    "\n",
    "        s.sort()\n",
    "        mapping = {s[i]:i for i in range(len(s))}\n",
    "        n = len(s)\n",
    "        seg_tree = [0] * (4*n + 1)\n",
    "        def query(lo, hi, node, left, right):\n",
    "            # print(lo, hi, node, left, right)\n",
    "            if lo <= left <= right <= hi:\n",
    "                return seg_tree[node]\n",
    "            mid = (left + right) >> 1\n",
    "            if hi <= mid:\n",
    "                return query(lo, hi, 2 * node + 1, left, mid)\n",
    "            if lo > mid:\n",
    "                return query(lo, hi, 2 * node + 2, mid + 1, right)\n",
    "            return max(query(lo, mid, 2 * node + 1, left, mid), \n",
    "                        query(mid + 1, hi, 2 * node + 2, mid + 1, right))\n",
    "\n",
    "        def put(val, idx, node, left, right):\n",
    "            if left == right == idx:\n",
    "                seg_tree[node] = val\n",
    "                return\n",
    "            mid = (left + right) >> 1\n",
    "            if idx <= mid:\n",
    "                put(val, idx, 2 * node + 1, left, mid)\n",
    "            else:\n",
    "                put(val, idx, 2 * node + 2, mid + 1, right)\n",
    "            seg_tree[node] = max(seg_tree[node], val)\n",
    "            return\n",
    "\n",
    "\n",
    "        ans = float('-inf')\n",
    "        for i in range(n):\n",
    "            key = mapping[nums[i] - i]\n",
    "            val = query(0, key, 0, 0, n - 1)\n",
    "            val += nums[i]\n",
    "            ans = max(ans, val)\n",
    "            put(val, key, 0, 0, n - 1)\n",
    "            \n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:\n",
    "        n = len(nums)\n",
    "        ls = [item - index for index, item in enumerate(nums)]\n",
    "        r = len(ls)\n",
    "        l = 0\n",
    "        st = [0] * (4 * n + 1)\n",
    "        mapV = {b:a for a, b in enumerate(sorted(ls))}\n",
    "\n",
    "        def update(index, val, node = 1, left = l, right= r):\n",
    "            if right < index or left > index: return\n",
    "            if left == right and left == index:\n",
    "                st[node] = val\n",
    "                return \n",
    "\n",
    "            mid = left + right >> 1\n",
    "            update(index,val,node * 2,left, mid)\n",
    "            update(index,val, node * 2 + 1, mid + 1, right)\n",
    "            st[node] = max(st[node * 2], st[node * 2 + 1])\n",
    "        \n",
    "      \n",
    "        def query(L, R, node = 1, left = l, right=r):\n",
    "            if right < L or left > R: return 0\n",
    "            if left >= L and right <= R: \n",
    "                return st[node]\n",
    "\n",
    "            mid = left + right >> 1\n",
    "            return max(\n",
    "                query(L, R, node * 2, left, mid),\n",
    "                query(L, R, node * 2 + 1, mid + 1, right),\n",
    "            )\n",
    "\n",
    "        res = -inf\n",
    "        for index, item in enumerate(nums):\n",
    "            d = item - index\n",
    "            t = query(l, mapV[d])\n",
    "            update(mapV[d],  t + max(item, 0))\n",
    "            res = max(res, t + item)\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:\n",
    "        n = len(nums)\n",
    "        f = nums[:]\n",
    "\n",
    "        tr = [-inf] * (4*n+4)\n",
    "        mv = {v:i for i,v in enumerate(sorted(set(nums[i]-i for i in range(n))))}\n",
    "\n",
    "        def modify(o:int,l:int,r:int,x:int,v:int)->None:\n",
    "            if l==r:\n",
    "                tr[o]=max(tr[o],v)\n",
    "                return\n",
    "            m = (l+r)//2\n",
    "            if x<=m:\n",
    "                modify(o<<1,l,m,x,v)\n",
    "            else:\n",
    "                modify(o<<1|1,m+1,r,x,v)\n",
    "            tr[o]=max(tr[o<<1],tr[o<<1|1])\n",
    "\n",
    "        def query(o:int,l:int,r:int,L:int,R:int)->int:\n",
    "            if L<=l and r<=R:\n",
    "                return tr[o]\n",
    "            m = (l+r)//2\n",
    "            if R<=m:\n",
    "                return query(o<<1,l,m,L,R)\n",
    "            if L>m:\n",
    "                return query(o<<1|1,m+1,r,L,R)\n",
    "            return max(query(o<<1,l,m,L,R),query(o<<1|1,m+1,r,L,R))\n",
    "\n",
    "\n",
    "        for i in range(n):\n",
    "            # for j in range(i):\n",
    "            #     if nums[i]-nums[j] >= i-j:\n",
    "            #         f[i]=max(f[i],f[j]+nums[i])\n",
    "            '''\n",
    "                (nums[i]-i) - (nums[j]-j) >=0\n",
    "                离散化做法\n",
    "            '''\n",
    "            f[i]=max(f[i],query(1,0,len(mv)-1,0,mv[nums[i]-i])+nums[i])\n",
    "            modify(1,0,len(mv)-1,mv[nums[i]-i],f[i])\n",
    "\n",
    "        return max(f)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# from sortedcontainers import SortedList\n",
    "# class Solution:\n",
    "#     def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:\n",
    "#         ans = max(nums)\n",
    "#         sl = SortedList([(-inf, 0)])\n",
    "#         for i, x in enumerate(nums):\n",
    "#             if x <= 0: continue\n",
    "#             p = sl.bisect_left((x - i + 1, 0))\n",
    "#             cur = sl[p - 1][1] + x\n",
    "#             ans = max(ans, cur)\n",
    "#             while p < len(sl) and sl[p][1] <= cur:\n",
    "#                 sl.pop(p)\n",
    "#             sl.add((x - i, cur))\n",
    "#         return ans\n",
    "\n",
    "class Solution:\n",
    "    def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:\n",
    "        n = len(nums)\n",
    "        f = [0] * (4 * n)\n",
    "\n",
    "        def update(k, l, r, x, v):\n",
    "            if l == r:\n",
    "                f[k] = v\n",
    "                return\n",
    "            m = l+r>>1\n",
    "            if x <= m: update(k<<1, l, m, x, v)\n",
    "            else: update(k<<1|1, m+1, r, x, v)\n",
    "            f[k] = max(f[k<<1], f[k<<1|1])\n",
    "\n",
    "        def query(k, l, r, L, R):\n",
    "            if L > R: return 0\n",
    "            if l == L and r == R: return f[k]\n",
    "            m = l+r>>1\n",
    "            return max(query(k<<1, l, m, L, min(R, m)), query(k<<1|1, m+1, r, max(L, m+1), R))\n",
    "\n",
    "        vals = sorted(nums[i] - i for i in range(n))\n",
    "        mp = {vals[i] : i for i in range(n)}\n",
    "\n",
    "        ans = max(nums)\n",
    "        for i in range(n):\n",
    "            pre = max(0, query(1, 1, n, 1, mp[nums[i] - i] + 1))\n",
    "            cur = pre + nums[i]\n",
    "            ans = max(ans, cur)\n",
    "            update(1, 1, n, mp[nums[i] - i] + 1, cur)\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "import numpy as np\n",
    "class Solution:\n",
    "    def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:\n",
    "        n = len(nums)\n",
    "        nums2 = []\n",
    "        for i,num in enumerate(nums):\n",
    "            nums2.append(num-i)\n",
    "        idxs = sorted(range(n),key=lambda i:nums2[i])\n",
    "        for i,x in enumerate(idxs):\n",
    "            nums2[x] = i\n",
    "        mp = defaultdict(int)\n",
    "        for i,x in enumerate(nums2):\n",
    "            mp[x]=i\n",
    "\n",
    "\n",
    "        tree = [-inf]*(n+1)\n",
    "        def lowbit(x):\n",
    "            return x&-x\n",
    "        def update(pos,v):\n",
    "            while pos<=n:\n",
    "                tree[pos]=max(tree[pos],v)\n",
    "                pos+=lowbit(pos)\n",
    "        def query(pos):\n",
    "            ret = -inf\n",
    "            while pos:\n",
    "                ret=max(ret,tree[pos])\n",
    "                pos-=lowbit(pos)\n",
    "            return ret\n",
    "\n",
    "        for idx in nums2:\n",
    "            f = max(query(idx+1), 0)\n",
    "            update(idx+1,f+nums[mp[idx]])\n",
    "        return query(n)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "import sys\n",
    "\n",
    "\n",
    "class SegmentTree:\n",
    "    \"\"\"线段树\n",
    "    Application:\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, lt):\n",
    "        self._n = len(lt)\n",
    "        self._max = [0] * (4 * self._n + 1)\n",
    "\n",
    "        def init(k, p, q):\n",
    "            if p == q:\n",
    "                self._max[k] = lt[p]\n",
    "                return\n",
    "            m = (p + q) >> 1\n",
    "            init(k << 1, p, m)\n",
    "            init((k << 1) + 1, m + 1, q)\n",
    "            self._max[k] = max(self._max[k << 1], self._max[(k << 1) + 1])\n",
    "\n",
    "        init(1, 0, self._n - 1)\n",
    "\n",
    "    def update(self, idx, val):\n",
    "        def dfs(k, p, q):\n",
    "            if p == q:\n",
    "                self._max[k] = val\n",
    "                return\n",
    "            m = (p + q) >> 1\n",
    "            if idx <= m:\n",
    "                dfs(k << 1, p, m)\n",
    "            else:\n",
    "                dfs((k << 1) + 1, m + 1, q)\n",
    "            self._max[k] = (\n",
    "                max(self._max[k << 1], self._max[(k << 1) + 1])\n",
    "            )\n",
    "\n",
    "        dfs(1, 0, self._n - 1)\n",
    "\n",
    "\n",
    "    def get_range_max(self, u, v):\n",
    "        def dfs(k, p, q, u, v):\n",
    "            if p == u and q == v:\n",
    "                return self._max[k]\n",
    "            m = (p + q) >> 1\n",
    "            ans = -sys.maxsize\n",
    "            if u <= m:\n",
    "                ans = max(ans, dfs(k << 1, p, m, u, min(m, v)))\n",
    "            if v > m:\n",
    "                ans = max(ans, dfs((k << 1) + 1, m + 1, q, max(m + 1, u), v))\n",
    "            return ans\n",
    "\n",
    "        return dfs(1, 0, self._n - 1, u, v)\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def maxBalancedSubsequenceSum(self, s: List[int]) -> int:\n",
    "        t = sorted(list(set([v - i for i, v in enumerate(s)])))\n",
    "        mp = {v: k for k, v in enumerate(t)}\n",
    "        n = len(t)\n",
    "        \n",
    "        st = SegmentTree([0] * n)\n",
    "        \n",
    "        ans = max(s)\n",
    "        for i in range(len(s)):\n",
    "            d = st.get_range_max(0, mp[s[i] - i])\n",
    "            st.update(mp[s[i] - i], d + s[i])\n",
    "            ans = max(ans, d + s[i])\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class SegmentTree:\n",
    "    def __init__(self):\n",
    "        self.tree = defaultdict(lambda: float('-inf'))\n",
    "\n",
    "    def update(self, lo, hi, left, right, idx, val):\n",
    "        if hi < left or right < lo:\n",
    "            return\n",
    "        if left <= lo and hi <= right:\n",
    "            self.tree[idx] = val\n",
    "            return\n",
    "        mi = (lo + hi) // 2\n",
    "        self.update(lo, mi, left, right, idx*2, val)\n",
    "        self.update(mi+1, hi, left, right, idx*2+1, val)\n",
    "        self.tree[idx] = max(self.tree[idx*2], self.tree[idx*2+1])\n",
    "    \n",
    "    def query(self, lo, hi, left, right, idx):\n",
    "        if lo > right or hi < left:\n",
    "            return float('-inf')\n",
    "        if left <= lo and hi <= right:\n",
    "            return self.tree[idx]\n",
    "        \n",
    "        mi = (lo + hi) // 2\n",
    "        left_max = self.query(lo, mi, left, right, idx*2)\n",
    "        right_max = self.query(mi + 1, hi, left, right, idx*2+1)\n",
    "        return max(left_max, right_max)\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:\n",
    "        res = max(nums)\n",
    "        n = len(nums)\n",
    "        dp = nums[:]\n",
    "        nums_v2 = [v-i for i, v in enumerate(nums)]\n",
    "        maps = {v:i for i, v in enumerate(sorted(nums_v2))}\n",
    "        lo, hi = 0, len(nums_v2)\n",
    "        tree = SegmentTree()\n",
    "        for i in range(n):\n",
    "            v = max(0, tree.query(lo, hi, lo, maps[nums_v2[i]], 1))\n",
    "            res = max(res, v + nums[i])\n",
    "            tree.update(lo, hi, maps[nums_v2[i]], maps[nums_v2[i]], 1, v+nums[i])\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:\n",
    "        \"\"\"\n",
    "            nums[ij] - nums[ij-1] >= ij - ij-1\n",
    "            元素差 大于等于 下标差\n",
    "            \n",
    "            \n",
    "            nums[ij] - ij >= nums[ij-1] - ij-1\n",
    "            b[ij] >= b[ij-1] 考虑b数组 找里面的最长递增子序列\n",
    "            \n",
    "            最大平衡子序列和\n",
    "            \n",
    "            平衡子序列 等价于b 的一个 IS\n",
    "            求最大和 \n",
    "        \"\"\"\n",
    "        n = len(nums)\n",
    "        b = [(nums[i] - i, nums[i]) for i in range(n)]\n",
    "        \n",
    "        # 离散化\n",
    "        ss = set()\n",
    "        for x, _ in b:\n",
    "            ss.add(x)\n",
    "    \n",
    "        values = list(sorted(ss))\n",
    "        idxs = defaultdict(int)    \n",
    "        for idx, val in enumerate(values):\n",
    "            idxs[val] = idx + 1\n",
    "                \n",
    "        u = len(values) + 1\n",
    "        mx = [-inf] * (4 * u)\n",
    "        \n",
    "        def modify(o: int, l: int, r: int, i: int, val: int) -> None:\n",
    "            if l == r:\n",
    "                mx[o] = val\n",
    "                return\n",
    "            m = (l + r) // 2\n",
    "            if i <= m: modify(o * 2, l, m, i, val)\n",
    "            else: modify(o * 2 + 1, m + 1, r, i, val)\n",
    "            mx[o] = max(mx[o * 2], mx[o * 2 + 1])\n",
    "\n",
    "        # 返回区间 [L,R] 内的最大值\n",
    "        def query(o: int, l: int, r: int, L: int, R: int) -> int:\n",
    "            if L <= l and r <= R: return mx[o]\n",
    "            res = -inf\n",
    "            m = (l + r) // 2\n",
    "            if L <= m: res = query(o * 2, l, m, L, R)\n",
    "            if R > m: res = max(res, query(o * 2 + 1, m + 1, r, L, R))\n",
    "            return res\n",
    "        \n",
    "        # b的第一个维度满足单调不减\n",
    "        for r1, r2 in b:\n",
    "            k = idxs[r1]\n",
    "            # print(k, r2)\n",
    "            # [1,k]区间上最大值 + 取到的值\n",
    "            res = r2 + max(0, query(1, 1, u, 1, k))\n",
    "            modify(1, 1, u, k, res)\n",
    "        \n",
    "        return mx[1]\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:\n",
    "        n = len(nums)\n",
    "        cnt1 = [0] * (4 * n)\n",
    "        todo = [0] * (4 * n)\n",
    "\n",
    "        # 维护区间 1 的个数\n",
    "        def maintain(o: int) -> None:\n",
    "            cnt1[o] = max(cnt1[o * 2], cnt1[o * 2 + 1])\n",
    "\n",
    "        # 执行区间反转\n",
    "        def do(o: int, l: int, r: int, val) -> None:\n",
    "            cnt1[o] = max(val, cnt1[o])\n",
    "            todo[o] = val\n",
    "\n",
    "        # 初始化线段树   o,l,r=1,1,n\n",
    "        def build(o: int, l: int, r: int) -> None:\n",
    "            if l == r:\n",
    "                cnt1[o] = -inf\n",
    "                return\n",
    "            m = (l + r) // 2\n",
    "            build(o * 2, l, m)\n",
    "            build(o * 2 + 1, m + 1, r)\n",
    "            maintain(o)\n",
    "\n",
    "        # 更新前缀 [L,R]   o,l,r=1,1,n\n",
    "        def update(o: int, l: int, r: int, L: int, R: int, val) -> None:\n",
    "            if L <= l and r <= R:\n",
    "                do(o, l, r, val)\n",
    "                return\n",
    "            m = (l + r) // 2\n",
    "            if todo[o]:\n",
    "                do(o * 2, l, m, todo[0])\n",
    "                do(o * 2 + 1, m + 1, r, todo[0])\n",
    "                todo[o] = 0\n",
    "            if m >= L: update(o * 2, l, m, L, R, val)\n",
    "            if m < R: update(o * 2 + 1, m + 1, r, L, R, val)\n",
    "            maintain(o)\n",
    "        \n",
    "        def query(o: int, l: int, r: int, L: int, R: int) -> None:\n",
    "            if L <= l and r <= R:\n",
    "                return cnt1[o]\n",
    "            m = (l + r) // 2\n",
    "            if todo[o]:\n",
    "                do(o * 2, l, m, todo[0])\n",
    "                do(o * 2 + 1, m + 1, r, todo[0])\n",
    "                todo[o] = 0\n",
    "            left = right = -inf\n",
    "            if m >= L: left = query(o * 2, l, m, L, R)\n",
    "            if m < R: right = query(o * 2 + 1, m + 1, r, L, R)\n",
    "            return max(left, right)\n",
    "    \n",
    "\n",
    "        build(1, 1, n)\n",
    "        b = sorted([x - i for i, x in enumerate(nums) if x > 0])\n",
    "        b2idx = {x : (i + 1) for i, x in enumerate(b)}\n",
    "        print(b2idx)\n",
    "        ans = max(nums)\n",
    "        for i, x in enumerate(nums):\n",
    "            if x <= 0:\n",
    "                continue\n",
    "            diff = x - i\n",
    "            idx = bisect_left(b, diff) + 1\n",
    "            dp_max = query(1, 1, n, 1, idx)\n",
    "            cur = max(dp_max + x, x)\n",
    "            update(1, 1, n, idx, idx, cur)\n",
    "            ans = max(ans, cur)\n",
    "        return ans\n",
    "            \n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:\n",
    "        n = len(nums)\n",
    "        nn = [x-i for i, x in enumerate(nums)]\n",
    "        idx = {x:i for i, x in zip(range(n), sorted(nn))}\n",
    "        g = [idx[x] for x in nn]\n",
    "    \n",
    "        MAX = max(g) + 1\n",
    "        rt = SegmentTree().root\n",
    "        for i, x in enumerate(g):\n",
    "            mx = SegmentTree.query(rt, 0, MAX, 0, x)\n",
    "            SegmentTree.update(rt, 0, MAX, x, x, max(mx, 0)+nums[i])\n",
    "        return SegmentTree.query(rt, 0, MAX, 0, MAX)\n",
    "        \n",
    "class Node:\n",
    "    def __init__(self) -> None:\n",
    "        self.ls = self.rs = None\n",
    "        self.val = -inf\n",
    "        self.add = 0\n",
    "\n",
    "class SegmentTree:\n",
    "    def __init__(self):\n",
    "        self.root = Node()\n",
    "    \n",
    "    @staticmethod\n",
    "    def update(node, lc, rc, l, r, v):\n",
    "        if l <= lc and rc <= r:\n",
    "            node.add = v\n",
    "            node.val = v\n",
    "            return\n",
    "        SegmentTree.pushdown(node)\n",
    "        mid = (lc + rc) // 2\n",
    "        if l <= mid:\n",
    "            SegmentTree.update(node.ls, lc, mid, l, r, v)\n",
    "        if r > mid:\n",
    "            SegmentTree.update(node.rs, mid + 1, rc, l, r, v)\n",
    "        node.val = max(node.ls.val, node.rs.val)\n",
    " \n",
    "    @staticmethod\n",
    "    def query(node, lc, rc, l, r):\n",
    "        if l <= lc and rc <= r:\n",
    "            return node.val\n",
    "        SegmentTree.pushdown(node)\n",
    "        mid, ans = (lc + rc) // 2, 0\n",
    "        if l <= mid:\n",
    "            ans = SegmentTree.query(node.ls, lc, mid, l, r)\n",
    "        if r > mid:\n",
    "            ans = max(ans, SegmentTree.query(node.rs, mid+1, rc, l, r))\n",
    "        return ans\n",
    "    \n",
    "    @staticmethod\n",
    "    def pushdown(node):\n",
    "        if node.ls is None: node.ls = Node()\n",
    "        if node.rs is None: node.rs = Node()\n",
    "        if not node.add: return\n",
    "        node.ls.add = node.add\n",
    "        node.rs.add = node.add\n",
    "        node.ls.val = node.add\n",
    "        node.rs.val = node.add\n",
    "        node.add = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# 动态区间 + 延迟更新 + 最值\n",
    "class DynamicSegmentTree:\n",
    "    def __init__(self, l, r, mx = 0):\n",
    "        self.l = l\n",
    "        self.r = r\n",
    "        self.mx = mx\n",
    "        self.left = None\n",
    "        self.right = None\n",
    "        self.todo = 0\n",
    "\n",
    "    def __maintain(self):\n",
    "        self.mx = max(self.left.mx, self.right.mx)\n",
    "\n",
    "    def __do(self):\n",
    "        self.left.mx += self.todo\n",
    "        self.left.todo += self.todo\n",
    "        self.right.mx += self.todo\n",
    "        self.right.todo += self.todo\n",
    "\n",
    "    def __check(self):\n",
    "        if self.left is None:\n",
    "            m = (self.l + self.r) // 2\n",
    "            self.left = DynamicSegmentTree(self.l, m, self.mx)\n",
    "            self.right = DynamicSegmentTree(m+1, self.r, self.mx)\n",
    "            self.todo = 0  # 创建新区间时，需要清除lazy标记\n",
    "\n",
    "    def add(self, L, R, v):\n",
    "        if L <= self.l and self.r <= R:\n",
    "            self.mx += v\n",
    "            self.todo += v\n",
    "            return\n",
    "\n",
    "        self.__check()\n",
    "        \n",
    "        if self.todo:\n",
    "            # do\n",
    "            self.__do()\n",
    "            self.todo = 0\n",
    "        m = (self.l + self.r) // 2\n",
    "        if m >= L:\n",
    "            self.left.add(L, R, v)\n",
    "        if m + 1 <= R:\n",
    "            self.right.add(L, R, v)\n",
    "        self.__maintain()\n",
    "\n",
    "    def query(self, L, R):\n",
    "        if L <= self.l and self.r <= R:\n",
    "            return self.mx\n",
    "        \n",
    "        self.__check()\n",
    "        \n",
    "        if self.left is None:\n",
    "            return self.mx\n",
    "        \n",
    "        m = (self.l + self.r) // 2\n",
    "        if self.todo:\n",
    "            # do\n",
    "            self.__do()\n",
    "            self.todo = 0\n",
    "        ans = 0\n",
    "        # m = (self.l + self.r) // 2\n",
    "        if m >= L:\n",
    "            ans = self.left.query(L, R)\n",
    "        if m + 1 <= R:\n",
    "            ans = max(ans, self.right.query(L, R))\n",
    "        return ans\n",
    "    \n",
    "class Solution:\n",
    "    def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:\n",
    "        ans = max(nums)\n",
    "        s = set()\n",
    "        for i, v in enumerate(nums):\n",
    "            if v < 0:\n",
    "                continue\n",
    "            s.add(v - i)\n",
    "        indexs = {}\n",
    "        for i, v in enumerate(sorted(s)):\n",
    "            indexs[v] = i + 1\n",
    "        q = DynamicSegmentTree(1, len(s))\n",
    "        \n",
    "        def find(l):\n",
    "            v = q.query(1, l)\n",
    "            return v\n",
    "        \n",
    "        def add(index, v):\n",
    "            raw = q.query(index, index)\n",
    "            #print(index, raw, v - raw)\n",
    "            q.add(index, index,  v - raw)\n",
    "            #print(\"after add\", index, v, v - raw, q.query(index, index))\n",
    "            \n",
    "        for i, v in enumerate(nums):\n",
    "            if v < 0:\n",
    "                continue\n",
    "            j = indexs[v-i]\n",
    "            t = find(j)\n",
    "            nv = max(t + v, v)\n",
    "            #print(i, v, v - i, t, nv)\n",
    "            add(j, nv)\n",
    "            ans = max(ans, nv)\n",
    "        \n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Node:\n",
    "    def __init__(self) -> None:\n",
    "        self.ls = self.rs = None\n",
    "        self.val = self.add = -2000000000\n",
    "\n",
    "class SegmentTree:\n",
    "    def __init__(self):\n",
    "        self.root = Node()\n",
    "\n",
    "    @staticmethod\n",
    "    def update(node: Node, lc: int, rc: int, l: int, r: int, v: int) -> None:\n",
    "        if l <= lc and rc <= r:\n",
    "            node.add = v\n",
    "            node.val = v\n",
    "            return\n",
    "        SegmentTree.pushdown(node)\n",
    "        mid = (lc + rc) >> 1\n",
    "        if l <= mid:\n",
    "            SegmentTree.update(node.ls, lc, mid, l, r, v)\n",
    "        if r > mid:\n",
    "            SegmentTree.update(node.rs, mid + 1, rc, l, r, v)\n",
    "        SegmentTree.pushup(node)\n",
    "\n",
    "    @staticmethod\n",
    "    def query(node: Node, lc: int, rc: int, l: int, r: int) -> int:\n",
    "        if l <= lc and rc <= r:\n",
    "            return node.val\n",
    "        SegmentTree.pushdown(node)\n",
    "        mid, ans = (lc + rc) >> 1, 0\n",
    "        if l <= mid:\n",
    "            ans = SegmentTree.query(node.ls, lc, mid, l, r)\n",
    "        if r > mid:\n",
    "            temp = SegmentTree.query(node.rs, mid + 1, rc, l, r)\n",
    "            if temp > ans:\n",
    "                ans = temp\n",
    "        return ans\n",
    "\n",
    "    @staticmethod\n",
    "    def pushdown(node: Node) -> None:\n",
    "        if node.ls is None:\n",
    "            node.ls = Node()\n",
    "        if node.rs is None:\n",
    "            node.rs = Node()\n",
    "        if not node.add:\n",
    "            return\n",
    "        node.ls.add = node.add\n",
    "        node.rs.add = node.add\n",
    "        node.ls.val = node.add\n",
    "        node.rs.val = node.add\n",
    "        node.add = 0\n",
    "\n",
    "    @staticmethod\n",
    "    def pushup(node: Node) -> None:\n",
    "        if node.ls.val > node.rs.val:\n",
    "            node.val = node.ls.val\n",
    "        else:\n",
    "            node.val = node.rs.val\n",
    "\n",
    "class Solution:\n",
    "    def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:\n",
    "        st = SegmentTree()\n",
    "        s = {num - index for index, num in enumerate(nums)}\n",
    "        d = {num: index for index, num in enumerate(sorted(s))}\n",
    "        max_range = len(d)\n",
    "        ans = nums[0]\n",
    "        for i, num in enumerate(nums):\n",
    "            index = d[num - i]\n",
    "            pre = st.query(st.root, 0, max_range, 0, index)\n",
    "            cur = num\n",
    "            if pre > 0:\n",
    "                cur += pre\n",
    "            if cur > ans:\n",
    "                ans = cur\n",
    "            st.update(st.root, 0, max_range, index, index, cur)\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "'''\n",
    "线段树\n",
    "'''\n",
    "\n",
    "\n",
    "# 线段树的节点类\n",
    "class STNode:\n",
    "    def __init__(self, left, right, val):\n",
    "        self.left = left  # 区间左边界\n",
    "        self.right = right  # 区间右边界\n",
    "        self.mid = (left + right) >> 1\n",
    "        self.leftNode = None  # 区间左节点\n",
    "        self.rightNode = None  # 区间右节点\n",
    "        self.val = val  # 节点值（区间值）\n",
    "        self.lazy_tag = None  # 区间问题的延迟更新标记\n",
    "\n",
    "    def __repr__(self):\n",
    "        return f\"{self.left}:{self.right} = {self.val}\"\n",
    "        \n",
    "from enum import Enum\n",
    "\n",
    "\n",
    "class SegOpCollect(Enum):\n",
    "    Collect_Sum = 1\n",
    "    Collect_Max = 2\n",
    "    Collect_Min = 3\n",
    "\n",
    "\n",
    "class SegOpUpdate(Enum):\n",
    "    Update_SetVal = 1\n",
    "    Update_AddVal = 2\n",
    "    Update_MultiplyVal = 3\n",
    "    Update_Reverse = 4  # 区间反转。refer to: https://leetcode.cn/problems/handling-sum-queries-after-update/solution/xian-duan-shu-by-tsreaper-2qot/\n",
    "\n",
    "\n",
    "# 线段树类\n",
    "class SegmentTree:\n",
    "    def __init__(self, collect=SegOpCollect.Collect_Max, update=SegOpUpdate.Update_SetVal, default_val=0):\n",
    "        self.op_collect = collect\n",
    "        self.op_update = update\n",
    "        self.default_val = default_val\n",
    "\n",
    "    def build(self, left=0, right=int(1e9)):\n",
    "        # 注意right不能小于查询边界，这里更新查询最大只能[left,right]\n",
    "        if self.op_collect == SegOpCollect.Collect_Sum:\n",
    "            self.root = STNode(left, right, self.default_val * (right - left + 1))\n",
    "        else:\n",
    "            self.root = STNode(left, right, self.default_val)\n",
    "\n",
    "    def build_with_data(self, data=[]):\n",
    "        # 注意right不能小于查询边界，这里更新查询最大只能[left,right]\n",
    "        def create_tree(left, right) -> STNode:  # 构建[left, right]的树\n",
    "            node = STNode(left, right, self.default_val)\n",
    "            if left == right:\n",
    "                node.val = data[left]  # 到达叶节点\n",
    "            else:\n",
    "                mid = (left + right) >> 1\n",
    "                # 构建左右子树\n",
    "                node.leftNode, node.rightNode = create_tree(left, mid), create_tree(mid + 1, right)\n",
    "                self.__merge(node)\n",
    "            return node\n",
    "\n",
    "        self.root = create_tree(0, len(data) - 1)\n",
    "\n",
    "    # collect 左区间和右区间的聚合方法；传入左区间和右区间，返回一个两区间的聚合结果\n",
    "    def collect(self, a, b):\n",
    "        if self.op_collect == SegOpCollect.Collect_Sum:\n",
    "            # 取sum\n",
    "            return a + b\n",
    "        elif self.op_collect == SegOpCollect.Collect_Max:\n",
    "            # 取max\n",
    "            return max(a, b)\n",
    "        elif self.op_collect == SegOpCollect.Collect_Min:\n",
    "            # 取min\n",
    "            return min(a, b)\n",
    "\n",
    "    # 用于加lazy标记的，实际上是当前val变下，然后加的标记影响子节点\n",
    "    def lazy_update(self, node: STNode, lazy: int):\n",
    "        # 源头是update_interval时调用，然后pushdown时有lazy_tag就下推，如果没有update_interval就没有lazy_tag了\n",
    "        meta_size = (node.right - node.left + 1)  # (r - l + 1)为此区间的元线段数\n",
    "        diff_val = None\n",
    "        if self.op_collect == SegOpCollect.Collect_Sum:\n",
    "            # - when sum\n",
    "            diff_val = meta_size * lazy\n",
    "        # elif self.op_collect == SegOpCollect.Collect_Max or self.op_collect == SegOpCollect.Collect_Min:\n",
    "        else:\n",
    "            # - when max/min\n",
    "            diff_val = lazy\n",
    "\n",
    "        if self.op_update == SegOpUpdate.Update_SetVal:\n",
    "            # - when set val\n",
    "            # 不累加标记\n",
    "            node.lazy_tag = lazy\n",
    "            node.val = diff_val\n",
    "        elif self.op_update == SegOpUpdate.Update_AddVal:\n",
    "            # - when add\n",
    "            # 累加标记\n",
    "            node.lazy_tag = node.lazy_tag + lazy if node.lazy_tag is not None else lazy\n",
    "            node.val += diff_val\n",
    "        elif self.op_update == SegOpUpdate.Update_MultiplyVal:\n",
    "            # - when multiply\n",
    "            # 累乘标记\n",
    "            node.lazy_tag = node.lazy_tag * lazy if node.lazy_tag is not None else lazy\n",
    "            node.val *= diff_val\n",
    "        elif self.op_update == SegOpUpdate.Update_Reverse:\n",
    "            # - when reverse, 必须是个sum树\n",
    "            node.lazy_tag = (node.lazy_tag + lazy) % 2 if node.lazy_tag is not None else lazy\n",
    "            # if node.lazy_tag == 0:\n",
    "            #     node.lazy_tag = None\n",
    "            if lazy == 1:\n",
    "                node.val = meta_size - node.val\n",
    "\n",
    "    def __merge(self, node):\n",
    "        node.val = self.collect(node.leftNode.val, node.rightNode.val)\n",
    "\n",
    "    # 向下更新 node 节点所在区间的左右子节点的值和懒惰标记, 实际是把当前节点的lazy flag去掉，转化并加到下面节点上去\n",
    "    def __pushdown(self, node):\n",
    "        # create child nodes\n",
    "        if not node.leftNode:\n",
    "            node.leftNode = STNode(node.left, node.mid, self.default_val)\n",
    "        if not node.rightNode:\n",
    "            node.rightNode = STNode(node.mid + 1, node.right, self.default_val)\n",
    "        # push down lazy tag\n",
    "        if node.lazy_tag is not None:\n",
    "            # 如果是叶子节点已经由父节点更新这个叶子节点值了\n",
    "            if node.left != node.right:\n",
    "                self.lazy_update(node.leftNode, node.lazy_tag)\n",
    "                self.lazy_update(node.rightNode, node.lazy_tag)\n",
    "            node.lazy_tag = None  # 更新当前节点的懒惰标记\n",
    "\n",
    "    # 单点更新，将 nums[i] 更改为 val\n",
    "    def set_point(self, i, val):\n",
    "        # 单点更新，将 nums[i] 更改为 val。node 节点的区间为 [node.left, node.right]\n",
    "        def dfs(i, val, node):\n",
    "            # 单点更新只改了个点的val，不需要考虑lazy flag\n",
    "            if node.left == node.right:\n",
    "                # self.lazy_update(node, val)\n",
    "                node.val = val\n",
    "                return\n",
    "            self.__pushdown(node)\n",
    "            if i <= node.mid:  # 在左子树中更新节点值\n",
    "                dfs(i, val, node.leftNode)\n",
    "            else:  # 在右子树中更新节点值\n",
    "                dfs(i, val, node.rightNode)\n",
    "            self.__merge(node)\n",
    "\n",
    "        dfs(i, val, self.root)\n",
    "\n",
    "    # 区间查询，查询区间为 [q_left, q_right] 的区间值\n",
    "    def query_interval(self, left, right):\n",
    "        def dfs(node):\n",
    "            # corner case\n",
    "            if left > right:\n",
    "                return self.default_val\n",
    "            if left <= node.left and node.right <= right:  # 区间就是想要的区间\n",
    "                return node.val  # 直接返回节点值\n",
    "            self.__pushdown(node)\n",
    "            if right <= node.mid:\n",
    "                return dfs(node.leftNode)\n",
    "            elif left > node.mid:\n",
    "                return dfs(node.rightNode)\n",
    "            res_left = dfs(node.leftNode)\n",
    "            res_right = dfs(node.rightNode)\n",
    "            return self.collect(res_left, res_right)  # 返回左右子树元素值的聚合计算结果\n",
    "\n",
    "        return dfs(self.root)\n",
    "\n",
    "    # 求[left,right]区间的大于等于val的第一个位置, 二分, 需要是max树, 如果没找到返回-1\n",
    "    def query_first_not_less_val(self, left, right, val):\n",
    "        assert self.op_collect == SegOpCollect.Collect_Max\n",
    "\n",
    "        # 区间查询，在线段树的 [left, right] 区间范围中搜索区间为 [left, right] 的区间值\n",
    "        def dfs(node):\n",
    "            if right < node.left or left > node.right:\n",
    "                return -1\n",
    "            if left <= node.left and node.right <= right:  # 区间就是想要的区间\n",
    "                if node.val < val:\n",
    "                    return -1\n",
    "            if node.left == node.right:\n",
    "                if node.val < val:\n",
    "                    return -1\n",
    "                return node.left\n",
    "            self.__pushdown(node)\n",
    "            res_left = -1\n",
    "            if left <= node.mid:  # 在左子树中查询\n",
    "                res_left = dfs(node.leftNode)\n",
    "                if res_left != -1:\n",
    "                    return res_left\n",
    "            res_right = -1\n",
    "            if right > node.mid:  # 在右子树中查询\n",
    "                res_right = dfs(node.rightNode)\n",
    "            return self.collect(res_left, res_right)  # 返回左右子树元素值的聚合计算结果\n",
    "\n",
    "        return dfs(self.root)\n",
    "\n",
    "    # 求[left,right]区间的大于等于val的最后一个位置, 二分, 需要是max树, 如果没找到返回-1\n",
    "    def query_last_not_less_val(self, left, right, val):\n",
    "        assert self.op_collect == SegOpCollect.Collect_Max\n",
    "\n",
    "        # 区间查询，在线段树的 [left, right] 区间范围中搜索区间为 [left, right] 的区间值\n",
    "        def dfs(node):\n",
    "            if right < node.left or left > node.right:\n",
    "                return -1\n",
    "            if left <= node.left and node.right <= right:  # 区间就是想要的区间\n",
    "                if node.val < val:\n",
    "                    return -1\n",
    "            if node.left == node.right:\n",
    "                if node.val < val:\n",
    "                    return -1\n",
    "                return node.left\n",
    "            self.__pushdown(node)\n",
    "            res_right = -1\n",
    "            if right > node.mid:  # 在右子树中查询\n",
    "                res_right = dfs(node.rightNode)\n",
    "                if res_right != -1:\n",
    "                    return res_right\n",
    "            res_left = -1\n",
    "            if left <= node.mid:  # 在左子树中查询\n",
    "                res_left = dfs(node.leftNode)\n",
    "            return self.collect(res_left, res_right)  # 返回左右子树元素值的聚合计算结果\n",
    "\n",
    "        return dfs(self.root)\n",
    "\n",
    "    # 区间更新，将区间[left, right]值进行变化val的处理\n",
    "    def update_interval(self, left, right, val):\n",
    "        # 区间更新\n",
    "        def dfs(node):\n",
    "            if left <= node.left and node.right <= right:\n",
    "                self.lazy_update(node, val)\n",
    "                return\n",
    "            self.__pushdown(node)  # 向下更新节点所在区间的左右子节点的值和懒惰标记\n",
    "            if left <= node.mid:  # 在左子树中更新区间值\n",
    "                dfs(node.leftNode)\n",
    "            if right > node.mid:  # 在右子树中更新区间值\n",
    "                dfs(node.rightNode)\n",
    "            self.__merge(node)\n",
    "\n",
    "        dfs(self.root)\n",
    "\n",
    "        \n",
    "        \n",
    "        \n",
    "MOD = int(1e9 + 7)\n",
    "INF = float('inf')\n",
    "import sortedcontainers\n",
    "import bisect\n",
    "import heapq\n",
    "\n",
    "class Solution:\n",
    "    def maxBalancedSubsequenceSum(self, nums: List[int]) -> int:\n",
    "        # len = 12\n",
    "        d = []\n",
    "        n = len(nums)\n",
    "        l = INF\n",
    "        r = float('-inf')\n",
    "        for i in range(n):\n",
    "            if nums[i] > 0:\n",
    "                # l = min(l, nums[i] - i)\n",
    "                # r = max(r, nums[i] - i)\n",
    "                d.append((nums[i] - i, nums[i]))\n",
    "        if not d:\n",
    "            return max(nums)\n",
    "        seg = SegmentTree(SegOpCollect.Collect_Max, SegOpUpdate.Update_SetVal, default_val=0)\n",
    "        diff2index = {}\n",
    "        x = sorted(set([v[0] for v in d]))\n",
    "        for v in x:\n",
    "            diff2index[v] = len(diff2index)\n",
    "        # print(d)\n",
    "        # print(diff2index)\n",
    "        l, r = 0, len(diff2index)\n",
    "        # len = 12\n",
    "        seg.build(left=l, right=r)\n",
    "        for v in d:\n",
    "            t = seg.query_interval(l, diff2index[v[0]])\n",
    "            seg.set_point(diff2index[v[0]], t + v[1])\n",
    "        return seg.query_interval(l, r)"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 2
}
